import numpy as np
import time
import ctypes

def profiler(func):
    def wrapper(*args, **kw):
        start = time.time()
        res = func(*args, **kw)
        end = time.time()
        print("Function `{}` running time: {:.6f}s".format(func.__name__, end-start))
        return res
    return wrapper

@profiler
def conv(x, w, stride, pad, b = None):
    x = np.pad(x, ((0,0), (0,0), (pad,pad), (pad,pad)), 'constant')
    H_new = int(1+(x.shape[2]-w.shape[2])/stride)
    W_new = int(1+(x.shape[3]-w.shape[3])/stride)
    out = np.zeros((x.shape[0], w.shape[0], H_new, W_new))
    for i in range(H_new):
        for j in range(W_new):
            x_window = x[:, :, i*stride:i*stride+w.shape[2], j*stride:j*stride+w.shape[3]]   # shape:(N, C, HH, WW)
            for k in range(x_window.shape[0]):
                for l in range(w.shape[0]):
                    out[k, l,  i, j] = np.sum(x_window[k] * w[l])
                if b != None:
                    out[k, :, i, j] += b
    return out

# @params: g is 3x3 matrix, which corresponds to weight kernels 
# @params: d is 4x4 matrix, which corresponds to the input tile whose output is 2x2
def wino_2x2_3x3_s1_2d_kernel(d, g):
    # Y = A_trans * (G*g*G_trans · B_trans*d*B) * A
    A = np.array([[1, 0], [1, 1], [1, -1], [0, -1]])
    A_trans = A.T
    G = np.array([[1, 0, 0], [0.5, 0.5, 0.5], [0.5, -0.5, 0.5], [0, 0, 1]])
    G_trans = G.T
    B = np.array([[1, 0, 0, 0], [0, 1, -1, 1], [-1, 1, 1, 0], [0, 0, 0, -1]])
    B_trans = B.T

    tmp = G.dot(g).dot(G_trans) * B_trans.dot(d).dot(B)
    Y = A_trans.dot(tmp).dot(A)
    return Y

@profiler
def conv_wino_3x3_s1(x, w, stride, pad, b=None):
    n, ic, ih, iw = x.shape
    oc, _, kh, kw = w.shape
    assert stride==1, "Only Support the case that stride is 1"
    assert kh==3 and kw==3, "Only Support the case that both kh and kw are 3"

    x = np.pad(x, ((0,0), (0,0), (pad,pad), (pad,pad)), 'constant')
    oh = int(1+(x.shape[2]-w.shape[2])/stride)
    ow = int(1+(x.shape[3]-w.shape[3])/stride)
    out = np.zeros((n, oc, oh, ow))
    for n_ in range(n):
        for oc_ in range(oc):
            for oh_ in range(0, oh, 2):
                for ow_ in range(0, ow, 2):
                    for ic_ in range(ic):
                        d = x[n_, ic_, oh_:oh_+4, ow_:ow_+4]
                        g = w[oc_, ic_, :, :]
                        out[n_, oc_, oh_:oh_+2, ow_:ow_+2] += wino_2x2_3x3_s1_2d_kernel(d, g)
    return out

if __name__ == '__main__':
    # n, ic, h, w = 16, 16, 32, 32
    # oc, _, kh, kw = 32, 16, 3, 3
    # # np.random.seed(0)
    # x = np.random.rand(n, ic, h, w)
    # w = np.random.rand(oc, ic, kh, kw)
    # wino_res = conv_wino_3x3_s1(x, w, 1, 1, b=None)
    # slide_res = conv(x, w, 1, 1, b=None)
    # print("relative error: ", np.max(np.abs((wino_res-slide_res)/slide_res)))
    # print("absolute error: ", np.max(np.abs(wino_res-slide_res)))

    # a = np.ones(shape=(6, 6))
    # a[0, :] = 0
    # a[:, 0] = 0
    # print(a)
    # b = np.array([[1.0000000000, 0.0000000000, -2.5000004768, -0.0000001192, 1.0000001192, 0.0000000000], 
    #               [0.0000000000, 0.9428091049, 1.3333333731, -0.4714044929, -0.6666667461, 0.0000000000], 
    #               [0.0000000000, -0.9428089857, 1.3333334923, 0.4714045525, -0.6666667461, 0.0000000000], 
    #               [0.0000000000, -0.1178511307, -0.0833333358, 0.2357022613, 0.1666666865, 0.0000000000], 
    #               [0.0000000000, 0.1178511307, -0.0833333507, -0.2357022911, 0.1666666865, 0.0000000000], 
    #               [0.0000000000, 0.9999998808, -0.0000000596, -2.5000000000, 0.0000000000, 1.0000000000]])
    # # print(a)
    # # print(b)
    # bt =b.T
    # # print(bt)
    # print(bt.dot(a.dot(b)))
